# 比较两个模型的结果
# 需要两个json文件，对每一个query图像，比较两个模型的top5 tp分数。
# 若分数差异很大抠出person 图像，保存到本地。
# 保存格式：为每个query 保存一个文件夹，
# 其中的图像名称为low-1.jpg，表示准确率较低模型的top1 图像
# high-2.jpg 表示准确率较高模型的top2图像。

import json
import os
from PIL import Image
from PIL import ImageDraw
import random
from tqdm import tqdm

import shutil
import numpy as np
class NpEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        else:
            return super(NpEncoder, self).default(obj)
# 读取两个json文件
def loadjson(fn1,fn2):
    fp1 = open(fn1,"r",encoding="utf-8")
    json_data1 = json.load(fp1)
    fp2 = open(fn2, "r", encoding="utf-8")
    json_data2 = json.load(fp2)
    return json_data1, json_data2

if __name__ =="__main__":
    fn_low = "/disk2t/fby/new_version_add_oldversion_modify/tools/prw_workspace/prw47/prw47results_1000.pkl_reid_result.json"
    fn_high = "/disk2t/fby/new_version_add_oldversion_modify/tools/prw_workspace/gc_net_shuffel/gc_net_shuffelepoch_24.pkl_reid_result.json"
    img_dataset_dir = "/disk2t/fby/new_version_add_oldversion_modify/PRW-v16.04.20/frames"

    shutil.rmtree("./diff_model_imgs")
    img_dst_dir = "./diff_model_imgs"

    json_data_low, json_data_high = loadjson(fn_low, fn_high)

    firstname = []
    if json_data_low["map"] > json_data_high["map"] :
        firstname = ["high", "low"]
    elif json_data_low["map"] < json_data_high["map"]:
        firstname = ["low", "high"]
    else:
        firstname = ["1", "2"]

    # 每个query的结果和gt值
    results_low = json_data_low["results"]
    results_high = json_data_high["results"]
    assert len(results_high) == len(results_low)

    diff_dict_top5 = []
    diff_dict_top10 = []
    for i in range(len(results_high)):
        current_result_low = results_low[i]
        current_result_high = results_high[i]
        assert current_result_low["probe_img"] == current_result_high["probe_img"]
        # 所使用的是top 10 如果gt没有10个那么就没必要看了
        if len(current_result_low['probe_gt']) < 10:
            continue
        if not current_result_high['gallery_top10_avg_tp'] > 0.6:
            continue
        delt_tp_top5 = current_result_high['gallery_top5_avg_tp'] - current_result_low['gallery_top5_avg_tp']
        delt_tp_top5 = round (delt_tp_top5, 2)
        delt_tp_top10 = current_result_high['gallery_top10_avg_tp'] - current_result_low['gallery_top10_avg_tp']
        delt_tp_top10 = round(delt_tp_top10, 2)
        # 前5个至少错一个
        if delt_tp_top5 > 0.15 :
            diff_dict_top5.append({"delt_tp_top5": delt_tp_top5,"delt_tp_top10": delt_tp_top10,
                                   'probe_img': current_result_low['probe_img'], 'probe_roi': current_result_low['probe_roi'],
                              "result_low":current_result_low['gallery'],"result_high":current_result_high['gallery']})

        if delt_tp_top10 > 0.1:
            diff_dict_top10.append({"delt_tp_top5": delt_tp_top5, "delt_tp_top10": delt_tp_top10,
                                    'probe_img': current_result_low['probe_img'], 'probe_roi': current_result_low['probe_roi'],
                              "result_low": current_result_low['gallery'], "result_high": current_result_high['gallery']})

    for diff in tqdm(diff_dict_top10):
        # 为每个query创建目录
        # 加随机数是防止重名的文件夹
        fn_dir = os.path.join(img_dst_dir, diff['probe_img'].split(".")[0] + "_" + str(random.randint(0, 9)))
        # assert os.path.exists(fn_dir) == False
        os.makedirs(fn_dir, exist_ok=True)

        # query 图像
        probe_img_dir = os.path.join(img_dataset_dir, diff['probe_img'])
        box = diff['probe_roi']
        img = Image.open(probe_img_dir)
        # draw = ImageDraw.ImageDraw(img)
        # draw.rectangle(((box[0], box[1]), (box[2], box[3])), fill=None, outline='yellow', width=5)
        crop = img.crop(box)
        crop.save(os.path.join(fn_dir, "probe_img.jpg"))

        jsonname = os.path.join(fn_dir, "model_diff.json")
        fjsonname = open(jsonname, "w", encoding="utf-8")
        json.dump(diff, fjsonname, ensure_ascii=False, cls=NpEncoder)

        for k, result in enumerate(diff["result_low"]):
            img_dir = os.path.join(img_dataset_dir, result["img"])
            box = result['roi'][:4:]
            img = Image.open(img_dir)
            crop = img.crop(box)
            crop = crop.resize((52, 128))
            if not result['correct'] == 1:
                draw = ImageDraw.ImageDraw(crop)
                draw.rectangle(((0, 0), (52, 128)), fill=None, outline='red', width=4)
            else:
                draw = ImageDraw.ImageDraw(crop)
                draw.rectangle(((0, 0), (52, 128)), fill=None, outline='green', width=4)
            fn = "low-top" + str(k + 1) + ".jpg"

            crop.save(os.path.join(fn_dir, fn))
        for k, result in enumerate(diff["result_high"]):
            img_dir = os.path.join(img_dataset_dir, result["img"])
            box = result['roi'][:4:]
            img = Image.open(img_dir)
            crop = img.crop(box)
            crop = crop.resize((52, 128))
            if not result['correct'] == 1:
                draw = ImageDraw.ImageDraw(crop)
                # draw.rectangle(((0, 0), (box[2]-box[0]-1, box[3]-box[1])), fill=None, outline='red', width=4)
                draw.rectangle(((0, 0), (52, 128)), fill=None, outline='red', width=4)
            else:
                draw = ImageDraw.ImageDraw(crop)
                draw.rectangle(((0, 0), (52, 128)), fill=None, outline='green', width=4)
            fn = "high-top" + str(k + 1) + ".jpg"

            crop.save(os.path.join(fn_dir, fn))
    a =10